#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Oct 18 21:16:53 2019

@author: ikaros
"""
import torch
torch.backends.cudnn.benchmark=True #网络的输入数据维度或类型上变化不大时可减少cudnn重复工作的次数，提高效率                 
import torch.nn as nn
import torchvision.models as models



class Model_1(nn.Module):
    def __init__(self):
        super().__init__()
        '''[resnet18,resnet34,resnet50,resnet101,resnet152]'''
        self.model = models.resnet50(pretrained=True)
        for param in self.model.parameters():
            param.requires_grad = False # 我们只是需要训练好的参数，不需要梯度信息，故设置为False
 
        #self.model.fc = nn.Linear(self.model.fc.in_features,2,bias=False) #修改全连接层
        #print(self.model.layer4) #查看网络结构
    def forward(self, x):
        x = self.model(x)
        return x

class Model_2(nn.Module):
    def __init__(self,model_path):
        super().__init__()
        '''[resnet18,resnet34,resnet50,resnet101,resnet152]'''
        self.model_path = model_path
        self.model = models.resnet50(pretrained=False)
        self.model.load_state_dict(torch.load(self.model_path))
        for param in self.model.parameters():
            param.requires_grad = False # 我们只是需要训练好的参数，不需要梯度信息，故设置为False
 
        self.model.fc = nn.Linear(self.model.fc.in_features,2,bias=False)
        '''括号里的是可以打印出来看的，并且可以通过调用它来修改'''
        #print(self.model)    #模型的全貌
        #print(self.model.fc) #全连接层
        #print(self.model.avgpool)  #平均池化
        #print(self.model.layer4)   
        #print(self.model.layer3)
        #print(self.model.layer2)
        #print(self.model.layer1)
        #print(self.model.layer1[0])
        #print(self.model.layer1[1])
        #print(self.model.layer1[2])
        #print(self.model.layer4[2].conv1)
        #print(self.model.layer4[2].bn1)
        #print(self.model.layer4[2].conv2)
        #print(self.model.layer4[2].bn2)
        #print(self.model.layer4[2].conv3)
        #print(self.model.layer4[2].bn3)
        #print(self.model.layer4[2].relu)
        '''修改结构'''
        #self.model.layer4[2].conv1 = nn.Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
        #print(self.model.layer4[2].conv1)
    def forward(self, x):
        x = self.model(x)
        return x

 
if __name__ == '__main__':
    model_path = './pths/resnet50-19c8e357.pth' #最好使用绝对路径
    model = Model_2(model_path)
    #print(model)
    
